package com.cloud.common.ribbon;

import com.alibaba.cloud.nacos.ribbon.NacosServer;
import com.cloud.common.context.AppContext;
import com.cloud.common.context.VersionContext;
import com.cloud.common.util.ReflectUtils;
import com.netflix.config.CachedDynamicIntProperty;
import com.netflix.loadbalancer.*;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.net.URL;
import java.util.*;
import java.util.stream.Collectors;

@Data
@Slf4j
public class RibbonBestRule extends ClientConfigEnabledRoundRobinRule {

    private static final String VERSION_KEY = "version";
    private static final ThreadLocal<String> chooseKeyHolder = new ThreadLocal<>();
    private static final ThreadLocal<String> serverHolder = new ThreadLocal<>();
    private static final Map<String, RibbonBestRule> ruleMap = new HashMap<>();
    private LoadBalancerStats loadBalancerStats;

    public static RibbonBestRule getInstance(String serviceName) {
        return ruleMap.get(serviceName);
    }

    public static void setChooseKey(String key) {
        chooseKeyHolder.set(key);
    }

    public static void removeChooseKey() {
        chooseKeyHolder.remove();
    }

    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        super.setLoadBalancer(lb);
        if (lb instanceof AbstractLoadBalancer) {
            loadBalancerStats = ((AbstractLoadBalancer) lb).getLoadBalancerStats();
            if (loadBalancerStats != null) {
                ruleMap.put(loadBalancerStats.getName(), this);
                try {
                    // 由于官方源码中这个stats对象的属性是私有的，而且默认值不太好，这里采用反射动态修改
                    ReflectUtils.setField(loadBalancerStats, "connectionFailureThreshold", new CachedDynamicIntProperty("ribbon.connectionFailureCountThreshold", 1));
                    ReflectUtils.setField(loadBalancerStats, "circuitTrippedTimeoutFactor", new CachedDynamicIntProperty("ribbon.circuitTripTimeoutFactorSeconds", 30));
                } catch (Exception e) {
                    log.error(e.getMessage(), e);
                }
            }
        }
    }

    @Override
    public Server choose(Object key) {
        List<Server> serverList = getLoadBalancer().getAllServers();
        if (CollectionUtils.isEmpty(serverList)) {
            return null;
        }

        // 从环境变量读取指定调用地址方便自测
        String requestURL = System.getenv(getServiceName());
        if (!StringUtils.isEmpty(requestURL)) {
            return buildServer(requestURL);
        }

        try {
            // 从线程本地变量获取指定服务器
            String forceServer = serverHolder.get();
            if (!StringUtils.isEmpty(forceServer)) {
                Optional<Server> optionalServer = serverList.stream()
                        .filter(server -> server.getId().equals(forceServer)).findFirst();
                if (optionalServer.isPresent()) {
                    return optionalServer.get();
                }
            }
        } finally {
            serverHolder.remove();
        }

        return doChoose();
    }

    /**
     * 选择服务器实例
     *
     * @return Server
     */
    public Server doChoose() {
        List<Server> serverList = getLoadBalancer().getAllServers();
        if (CollectionUtils.isEmpty(serverList)) {
            return null;
        }

        // 过滤调试的服务器并按版本号分组
        Map<String, List<NacosServer>> versionServerMap = serverList.stream().map(server -> (NacosServer) server)
                .filter(server -> !AppContext.ENV_LOCAL.equalsIgnoreCase(server.getMetadata().get("profile")))
                .collect(Collectors.groupingBy(server -> server.getMetadata().get(VERSION_KEY)));
        // 按版本号从高到低排序
        List<String> versionList = versionServerMap.keySet().stream()
                .sorted(Comparator.comparing(String::toString).reversed())
                .collect(Collectors.toList());
        // 默认获取低版本
        List<NacosServer> matchServerList = versionServerMap.get(versionList.get(versionList.size() - 1));

        // 根据上下文的版本号选择服务器
        String version = VersionContext.get();
        if (versionServerMap.size() > 1 && version != null) {
            for (String serverVersion : versionList) {
                if (serverVersion.compareTo(version) <= 0) {
                    matchServerList = versionServerMap.get(serverVersion);
                    break;
                }
            }
        }
        return bestChoose(matchServerList);
    }

    /**
     * 按最少请求数选择
     *
     * @param serverList
     * @return
     */
    private Server bestChoose(List<NacosServer> serverList) {
        if (loadBalancerStats == null || CollectionUtils.isEmpty(serverList)) {
            return null;
        }
        String key = chooseKeyHolder.get();
        if (key != null) {
            int index = Math.abs(key.hashCode()) % serverList.size();
            return serverList.get(index);
        }
        // 选择未被熔断并且连接次数最少
        Server chooseServer = null;
        long currentTime = System.currentTimeMillis();
        int minimalConcurrentConnections = Integer.MAX_VALUE;
        for (Server server : serverList) {
            ServerStats serverStats = loadBalancerStats.getSingleServerStat(server);
            if (serverStats.isCircuitBreakerTripped(currentTime)) {
                continue;
            }
            int concurrentConnections = serverStats.getActiveRequestsCount(currentTime);
            if (concurrentConnections < minimalConcurrentConnections) {
                minimalConcurrentConnections = concurrentConnections;
                chooseServer = server;
            }
        }
        return chooseServer != null ? chooseServer : super.choose(key);
    }

    /**
     * 指定服务器
     *
     * @param host
     * @param port
     */
    public void assignServer(String host, int port) {
        serverHolder.set(host + ":" + port);
    }

    private String getServiceName() {
        String serviceName = getLoadBalancer().getAllServers().stream().findFirst().get().getMetaInfo().getAppName().toLowerCase();
        // 兼容nacos
        int splitIndex = serviceName.indexOf("@@");
        if (splitIndex > 0) {
            serviceName = serviceName.substring(splitIndex + 2);
        }
        return serviceName;
    }

    private Server buildServer(String requestURL) {
        try {
            URL url = new URL(requestURL);
            return new Server(url.getProtocol(), url.getHost(), url.getPort());
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
        return null;
    }
}
